#include "vision/StdAfx.h"
#include "gmm.hpp"

namespace MT {

// -----------------------------------------------------------------------------------------
// distribution -- a distribution on a D-dimensional Euclidean space
// -----------------------------------------------------------------------------------------

// -----------------------------------------------------------------------------------------
// gaussian -- a multivariate gaussian distribution
// -----------------------------------------------------------------------------------------

// -----------------------------------------------------------------------------------------
// gmm -- a gaussian mixture model
// -----------------------------------------------------------------------------------------


// -----------------------------------------------------------------------------------------
// gaussian/gmm operations
// -----------------------------------------------------------------------------------------


gmm operator +(gaussian const &x, gaussian const &y)
{
    std::vector<gaussian const *> v(2); v[0] = &x; v[1] = &y;
    return gmm(v);
}

gmm operator +(gaussian const &x, gmm const &y)
{
    std::vector<gaussian const *> v;
    int i, n = y.g.size();
    v.resize(n+1);
    v[0] = &x;
    for(i = 0; i < n; ++i) v[i+1] = &y.g[i];
    return gmm(v);
}

gmm operator +(gmm const &x, gaussian const &y)
{
    std::vector<gaussian const *> v;
    int i, n = x.g.size();
    v.resize(n+1);
    for(i = 0; i < n; ++i) v[i] = &x.g[i];
    v[0] = &y;
    return gmm(v);
}

gmm operator +(gmm const &x, gmm const &y)
{
    std::vector<gaussian const *> v;
    int i, n1 = x.g.size(), n2 = y.g.size();
    v.resize(n1+n2);
    for(i = 0; i < n1; ++i) v[i] = &x.g[i];
    for(i = 0; i < n2; ++i) v[n1+i] = &y.g[i];
    return gmm(v);
}



gaussian operator *(gaussian const &x, gaussian const &y)
{
    CV_Assert(x.is_valid() && y.is_valid() && x.get_dim() == y.get_dim());
    cv::Mat c, C, invC;

    // compute covariance
    cv::add(x.get_invCov(), y.get_invCov(), invC);
    real detInvCov = cv::invert(invC, C);
    CV_Assert(detInvCov > 0);

    // Cov_x^{-1} mean_x
    cv::Mat x1; cv::gemm(x.get_mean(), x.get_invCov(), 1, cv::Mat(), 0, x1, cv::GEMM_2_T);

    // Cov_y^{-1} mean_y
    cv::Mat y1; cv::gemm(y.get_mean(), y.get_invCov(), 1, cv::Mat(), 0, y1, cv::GEMM_2_T);

    // compute mean
    cv::Mat z; cv::add(x1, y1, z);
    cv::gemm(z, C, 1, cv::Mat(), 0, c, cv::GEMM_2_T);

    // Cov_z^{-1} mean_z
    cv::gemm(c, invC, 1, cv::Mat(), 0, z, cv::GEMM_2_T);

    // compute the normalisation factor
    real w = -0.5*(x.get_mean().dot(x1) + y.get_mean().dot(y1) - c.dot(z));
    w = w < -100? 0: x.get_weight()*x.get_Z()*y.get_weight()*y.get_Z()*exp(w);

    // return the product gaussian
    gaussian g(c, C, 1.0, invC, 1.0/detInvCov);
    g.set_weight(w/g.get_Z());
    return g;
}

gmm operator *(gaussian const &x, gmm const &y)
{
    int n = y.g.size();
    std::vector<gaussian> g(n);
    while(--n >= 0) g[n] = x*y.g[n];
    return gmm(g);
}

gmm operator *(gmm const &x, gaussian const &y)
{
    int n = x.g.size();
    std::vector<gaussian> g(n);
    while(--n >= 0) g[n] = x.g[n]*y;
    return gmm(g);
}

gmm operator *(gmm const &x, gmm const &y)
{
    int n1 = x.g.size(), n2 = y.g.size();
    std::vector<gaussian> v(n1*n2);
    int i, j, cnt = 0;
    for(i = 0; i < n1; i++) for(j = 0; j < n2; ++j) v[cnt++] = x.g[i] * y.g[j];
    return gmm(v);
}




gaussian convolve(gaussian const &x, gaussian const &y)
{
    CV_Assert(x.is_valid() && y.is_valid() && x.get_dim() == y.get_dim());
    cv::Mat c, C;
    cv::add(x.get_mean(), y.get_mean(), c);
    cv::add(x.get_covariance(), y.get_covariance(), C);
    return gaussian(c, C, x.get_weight()*y.get_weight());
}

gmm convolve(gaussian const &x, gmm const &y)
{
    int n = y.g.size();
    std::vector<gaussian> g(n);
    while(--n >= 0) g[n] = convolve(x, y.g[n]);
    return gmm(g);
}

gmm convolve(gmm const &x, gaussian const &y)
{
    int n = x.g.size();
    std::vector<gaussian> g(n);
    while(--n >= 0) g[n] = convolve(x.g[n], y);
    return gmm(g);
}

gmm convolve(gmm const &x, gmm const &y)
{
    int n1 = x.g.size(), n2 = y.g.size();
    std::vector<gaussian> v(n1*n2);
    int i, j, cnt = 0;
    for(i = 0; i < n1; i++) for(j = 0; j < n2; ++j) v[cnt++] = convolve(x.g[i], y.g[j]);
    return gmm(v);
}


// reduce the number of gaussians by grouping clusters of gaussians
//
// Gaussians are grouped into clusters by running weighted K-Means clustering
// on top of their centers. Then, each cluster of gaussians is approximated
// by a single gaussian using maximum likelihood
gmm reduce(gmm const &inst, int K)
{
    CV_Assert(inst.is_valid());
    int n = inst.g.size();
    if(K >= n) return inst;

    // TODO: here, need a weighted K-Means clustering
    return gmm();
}

// reduce the gmm down to a single gaussian by using an unbiased estimation
gaussian reduce(gmm const &inst)
{
    CV_Assert(inst.is_valid());
    int n = inst.g.size();
    if(n==1)
    {
        gaussian const &g = inst.g[0];
        return gaussian(g.get_mean(), g.get_covariance(), inst.get_weight(), g.get_invCov(), g.get_detCov());
    }

    int D = inst.get_dim();
    cv::Mat mean = cv::Mat(1, D, CV_REAL, cv::Scalar()), cov = cv::Mat(D, D, CV_REAL, cv::Scalar());

    while(--n >= 0)
    {
        gaussian const &g = inst.g[n];
        cv::scaleAdd(g.get_mean(), g.get_weight(), mean, mean);
        cv::scaleAdd(g.get_covariance(), g.get_weight()*g.get_weight(), cov, cov);
    }

    return gaussian(mean, cov, inst.get_weight());
}

} // namespace MT

